{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Implementing a RAG System with Open Source SLMs and LanceDB\n",
        "This notebook is a companion of chapter 12 of the \"Small Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2025.  \n",
        "The code in this notebook is about implementing a basic RAG (Retrieval Augmented Generation) system using only Small Language Models (SLMs) and a Open Source vector database, [LanceDB](https://lancedb.github.io/lancedb/). The data preprocessing, embedding transformation, retrieval with or without the vector database, don't require hardware acceleration. The answer generation process can run properly with or without hardware acceleration, but to make it faster, loading the model weights to a GPU is recommended.   \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "CztxooKJTcWX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the missing requirements (only PyMuPDF missing in the Colab VM)."
      ],
      "metadata": {
        "id": "ZWaL4302HBDT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install PyMuPDF"
      ],
      "metadata": {
        "id": "PsYed2kiHGvm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data Preprocessing"
      ],
      "metadata": {
        "id": "1bsiWQDscj7y"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the knowledge base (a single PDF document)."
      ],
      "metadata": {
        "id": "CXA5pbGGHJtH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import requests\n",
        "\n",
        "pdf_path = \"2401.08671.pdf\"\n",
        "\n",
        "if not os.path.exists(pdf_path):\n",
        "  print(\"File doesn't exist, downloading it...\")\n",
        "\n",
        "  url = \"https://arxiv.org/pdf/2401.08671\"\n",
        "\n",
        "  filename = pdf_path\n",
        "\n",
        "  response = requests.get(url)\n",
        "\n",
        "  if response.status_code == 200:\n",
        "      with open(filename, \"wb\") as file:\n",
        "          file.write(response.content)\n",
        "      print(f\"The file has been downloaded and saved as {filename}\")\n",
        "  else:\n",
        "      print(f\"Failed to download the file. Status code: {response.status_code}\")\n",
        "else:\n",
        "  print(f\"File {pdf_path} exists.\")"
      ],
      "metadata": {
        "id": "bFCUKrw0HM-O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implement a custom function that uses PyMuPDF to parse the PDF document to extract the raw text only (no tables, nor images) and get some statistics about it. A separate function for text formatting has been implemented too, so that extra ad hoc formatting can be easily added any time without touching the main parse function."
      ],
      "metadata": {
        "id": "Xa-hUydIHQv7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import fitz\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "def text_formatter(text):\n",
        "    cleaned_text = text.replace(\"\\n\", \" \").strip()\n",
        "\n",
        "    # Add here any other extra text formatting\n",
        "\n",
        "    return cleaned_text\n",
        "\n",
        "def open_and_read_pdf(pdf_path):\n",
        "    \"\"\"\n",
        "    Opens a PDF file, reads its text content page by page, and collects statistics.\n",
        "\n",
        "    Parameters:\n",
        "        pdf_path (str): The file path to the PDF document to be opened and read.\n",
        "\n",
        "    Returns:\n",
        "        list[dict]: A list of dictionaries, each containing the page number\n",
        "        (adjusted), character count, word count, sentence count, token count, and the extracted text\n",
        "        for each page.\n",
        "    \"\"\"\n",
        "    doc = fitz.open(pdf_path)\n",
        "    pages_and_texts = []\n",
        "    for page_number, page in tqdm(enumerate(doc)):\n",
        "        text = page.get_text()\n",
        "        text = text_formatter(text)\n",
        "        pages_and_texts.append({\"page_number\": page_number + 1,\n",
        "                                \"page_char_count\": len(text),\n",
        "                                \"page_word_count\": len(text.split(\" \")),\n",
        "                                \"page_sentence_count_raw\": len(text.split(\". \")),\n",
        "                                \"page_token_count\": len(text) / 4,\n",
        "                                \"text\": text})\n",
        "    return pages_and_texts"
      ],
      "metadata": {
        "id": "Lmhq-R9TH8EQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Parse the PDF document."
      ],
      "metadata": {
        "id": "SgDOR9MqYYix"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pages_and_texts = open_and_read_pdf(pdf_path=pdf_path)\n",
        "pages_and_texts[:2]"
      ],
      "metadata": {
        "id": "kqb7-cDlYIxO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Split the document's text into sentences using [SpaCy](https://spacy.io/). Force each sentence to string format, just in case, and calculate also the count of sentences per page."
      ],
      "metadata": {
        "id": "bcIe4VqbH-Q6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from spacy.lang.en import English\n",
        "\n",
        "nlp = English()\n",
        "\n",
        "nlp.add_pipe(\"sentencizer\")"
      ],
      "metadata": {
        "id": "O6hd9S1PIRCJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for item in tqdm(pages_and_texts):\n",
        "    item[\"sentences\"] = list(nlp(item[\"text\"]).sents)\n",
        "\n",
        "    item[\"sentences\"] = [str(sentence) for sentence in item[\"sentences\"]]\n",
        "\n",
        "    item[\"page_sentence_count_spacy\"] = len(item[\"sentences\"])"
      ],
      "metadata": {
        "id": "cODlVr3_ISbO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to split a list into others of a given size. It will be used to split each page sentences into chunks."
      ],
      "metadata": {
        "id": "gYSqr6BpIVhz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_sentence_chunk_size = 10\n",
        "\n",
        "def split_list(input_list, slice_size):\n",
        "\n",
        "    return [input_list[i:i + slice_size] for i in range(0, len(input_list), slice_size)]"
      ],
      "metadata": {
        "id": "S6sSYCL_IoQe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Loop through pages and texts and split sentences into chunks."
      ],
      "metadata": {
        "id": "G2JbQQRXapOp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for item in tqdm(pages_and_texts):\n",
        "    item[\"sentence_chunks\"] = split_list(input_list=item[\"sentences\"],\n",
        "                                         slice_size=num_sentence_chunk_size)\n",
        "    item[\"num_chunks\"] = len(item[\"sentence_chunks\"])"
      ],
      "metadata": {
        "id": "EZNzYY3WaoPA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Split each of the created chunks into its own item (ad get also some stats about each chunk)."
      ],
      "metadata": {
        "id": "AoXDxJESOZFG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "\n",
        "pages_and_chunks = []\n",
        "for item in tqdm(pages_and_texts):\n",
        "    for sentence_chunk in item[\"sentence_chunks\"]:\n",
        "        chunk_dict = {}\n",
        "        chunk_dict[\"page_number\"] = item[\"page_number\"]\n",
        "\n",
        "        joined_sentence_chunk = \"\".join(sentence_chunk).replace(\"  \", \" \").strip()\n",
        "        joined_sentence_chunk = re.sub(r'\\.([A-Z])', r'. \\1', joined_sentence_chunk)\n",
        "        chunk_dict[\"sentence_chunk\"] = joined_sentence_chunk\n",
        "\n",
        "        chunk_dict[\"chunk_char_count\"] = len(joined_sentence_chunk)\n",
        "        chunk_dict[\"chunk_word_count\"] = len([word for word in joined_sentence_chunk.split(\" \")])\n",
        "        chunk_dict[\"chunk_token_count\"] = len(joined_sentence_chunk) / 4 # 1 token = ~4 characters\n",
        "\n",
        "        pages_and_chunks.append(chunk_dict)"
      ],
      "metadata": {
        "id": "-knU-I6MOclp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Calculate how many chucks we have."
      ],
      "metadata": {
        "id": "nmGKbgpVbGnr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "len(pages_and_chunks)"
      ],
      "metadata": {
        "id": "rhYE-1J3bIBQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Optional step: filter sentences that have less than 30 tokens, as they would probably contain negligible information."
      ],
      "metadata": {
        "id": "4gIr0LB3OBMI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "\n",
        "df = pd.DataFrame(pages_and_chunks)\n",
        "df.describe().round(2)"
      ],
      "metadata": {
        "id": "-3f2UKPeOPTQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "min_token_length = 30\n",
        "subset_df = df[df[\"chunk_token_count\"] <= min_token_length]\n",
        "\n",
        "for row in subset_df.sample(min(5, len(subset_df)), replace=True).iterrows():\n",
        "    print(f'Chunk token count: {row[1][\"chunk_token_count\"]} | Text: {row[1][\"sentence_chunk\"]}')"
      ],
      "metadata": {
        "id": "LKKkTREv3w0H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pages_and_chunks_over_min_token_len = df[df[\"chunk_token_count\"] > min_token_length].to_dict(orient=\"records\")\n",
        "pages_and_chunks_over_min_token_len[:2]"
      ],
      "metadata": {
        "id": "sx9_3MufOGC6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Select an embedding model ([all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2) in this example)."
      ],
      "metadata": {
        "id": "3sxX3sj-cr0O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "\n",
        "embedding_model = SentenceTransformer(model_name_or_path=\"all-mpnet-base-v2\",\n",
        "                                      device=\"cpu\")"
      ],
      "metadata": {
        "id": "4kZ9g-LOI2jb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Embed the chunks."
      ],
      "metadata": {
        "id": "8B7_d2u7IprD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for item in tqdm(pages_and_chunks_over_min_token_len):\n",
        "    item[\"embedding\"] = embedding_model.encode(item[\"sentence_chunk\"])"
      ],
      "metadata": {
        "id": "OWXnOf7hI_Iu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Save the embeddings to file."
      ],
      "metadata": {
        "id": "BDhe8A_UJRtm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "text_chunks_and_embeddings_df = pd.DataFrame(pages_and_chunks_over_min_token_len)\n",
        "text_chunks_and_embeddings_df['embedding_str'] = text_chunks_and_embeddings_df['embedding'].apply(\n",
        "    lambda x: np.array2string(x, separator=',')[1:-1])\n",
        "embeddings_df_save_path = \"text_chunks_and_embeddings_df.csv\"\n",
        "text_chunks_and_embeddings_df.to_csv(embeddings_df_save_path, index=False)"
      ],
      "metadata": {
        "id": "f9hRzcLWJToN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you want to skip the vector database integration, you can jump straight to the [Search (Embeddings on File)](#scrollTo=7G8J_IWRJYKP) section."
      ],
      "metadata": {
        "id": "s-KNZmTzdW4z"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### LanceDB"
      ],
      "metadata": {
        "id": "6RjcpGJimhv-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This section introduced a vector database, LanceDB, to store the document embeddings.  \n",
        "Let's install the LanceDB package first."
      ],
      "metadata": {
        "id": "jc_caJH2mkVK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install lancedb"
      ],
      "metadata": {
        "id": "BJkYy5jFxrCR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create a local empty database instance and estabilish a connection to it."
      ],
      "metadata": {
        "id": "mnevu60fm9MH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import lancedb\n",
        "\n",
        "db = lancedb.connect(\"paperdb\")"
      ],
      "metadata": {
        "id": "0i0UygD8mr2q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the preliminary saved embeddings from file and put them into a Pandas DataFrame."
      ],
      "metadata": {
        "id": "acMoLoSDnJNY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "text_chunks_and_embedding_df_load = pd.read_csv(embeddings_df_save_path)\n",
        "text_chunks_and_embedding_df_load['embedding_final'] = text_chunks_and_embedding_df_load['embedding_str'].apply(lambda x: np.fromstring(x, sep=','))\n",
        "text_chunks_and_embedding_df_load.head()"
      ],
      "metadata": {
        "id": "VevGVxYVnl4d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Prepare embeddings and other data for ingestion to the database."
      ],
      "metadata": {
        "id": "HQTSrdUzm9xH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "final_data = []\n",
        "for row in text_chunks_and_embedding_df_load.itertuples(index=False):\n",
        "    temp = {}\n",
        "    temp[\"page_number\"] = row.page_number\n",
        "    temp[\"chunk_char_count\"] = row.chunk_char_count\n",
        "    temp[\"chunk_word_count\"] = row.chunk_word_count\n",
        "    temp[\"chunk_token_count\"] = row.chunk_token_count\n",
        "    temp[\"sentence_chunk\"] = row.sentence_chunk\n",
        "    temp[\"embedding\"] = np.array(row.embedding_final)\n",
        "    final_data.append(temp)"
      ],
      "metadata": {
        "id": "RXB0pw2omtNv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Store embeddings and data in a LanceDB table ([PyArrow](https://arrow.apache.org/docs/python/index.html) format)."
      ],
      "metadata": {
        "id": "dbCfoELFuYeL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pyarrow as pa"
      ],
      "metadata": {
        "id": "RR-jVBnQua_2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_dict = {}\n",
        "for key in final_data[0].keys():\n",
        "    data_dict[key] = [d[key] for d in final_data]"
      ],
      "metadata": {
        "id": "xC5rI2NQPAjp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "table = db.create_table(\n",
        "    \"paper_embeddings_table\",\n",
        "    data=pa.Table.from_pydict(data_dict),\n",
        "    mode=\"overwrite\",\n",
        "    schema=pa.schema([\n",
        "        (\"page_number\", pa.int64()),\n",
        "        (\"chunk_char_count\", pa.int64()),\n",
        "        (\"chunk_word_count\", pa.int64()),\n",
        "        (\"chunk_token_count\", pa.float64()),\n",
        "        (\"sentence_chunk\", pa.string()),\n",
        "        (\"embedding\", pa.list_(pa.float32(), list_size=768)),\n",
        "    ]),\n",
        ")"
      ],
      "metadata": {
        "id": "Ii2KRyJpM0Do"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an index over the table."
      ],
      "metadata": {
        "id": "Pxh-58imvStK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "db[\"paper_embeddings_table\"].create_index(\n",
        "    metric=\"cosine\",\n",
        "    vector_column_name=\"embedding\",\n",
        "    index_type=\"IVF_FLAT\",\n",
        ")"
      ],
      "metadata": {
        "id": "Lr0GO6fp6jT5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Provide a query and convert it into embeddings."
      ],
      "metadata": {
        "id": "BLRNDCZtvrm4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"blocked KV-cache\"\n",
        "query_embedding = embedding_model.encode(query, convert_to_tensor=False)"
      ],
      "metadata": {
        "id": "Ek8_VUyZv1zJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Perform the search on the LanceDB table."
      ],
      "metadata": {
        "id": "8SnJ4V8cv8iE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "results = db[\"paper_embeddings_table\"].search(\n",
        "        query_embedding,\n",
        "        vector_column_name=\"embedding\"\n",
        "    ).limit(3).to_list()"
      ],
      "metadata": {
        "id": "kysYIu1HwD_u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Display the results."
      ],
      "metadata": {
        "id": "T8aDddJFwUD7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Search results:\")\n",
        "for result in results:\n",
        "    print(result['_distance'])\n",
        "    print(result['sentence_chunk'])\n",
        "    print(result['page_number'])\n",
        "    print(\"\")"
      ],
      "metadata": {
        "id": "SWIp9AYYwV0C"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Search (Embeddings on File)"
      ],
      "metadata": {
        "id": "7G8J_IWRJYKP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the embeddings preliminary saved on file and put them into a Pandas DataFrame."
      ],
      "metadata": {
        "id": "ptJcfDJNaUmD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "text_chunks_and_embedding_df = pd.read_csv(\"text_chunks_and_embeddings_df.csv\")"
      ],
      "metadata": {
        "id": "WqWmdtYHakOl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Convert the `embedding` column back to Numpy array (it was converted to string when saved to CSV)."
      ],
      "metadata": {
        "id": "wWIMaPcKhnDM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "text_chunks_and_embedding_df[\"embedding\"] = text_chunks_and_embedding_df[\"embedding\"].apply(lambda x: np.fromstring(x.strip(\"[]\"), sep=\" \"))"
      ],
      "metadata": {
        "id": "OVSgZoIDhjJy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Convert the texts and embedding DataFrame to a list of dictionaries and then convert the embeddings to Torch tensors (numeric format conversion is required, as the NumPy arrays are Float 64, while Torch tensors are Float 32 by default)."
      ],
      "metadata": {
        "id": "NLksdawchz8y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pages_and_chunks = text_chunks_and_embedding_df.to_dict(orient=\"records\")\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "embeddings = torch.tensor(np.array(text_chunks_and_embedding_df[\"embedding\"].tolist()), dtype=torch.float32).to(device)\n",
        "embeddings.shape"
      ],
      "metadata": {
        "id": "mEZxzM5VhzT4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create again an instance of the Sentence Transformer, jsut in case we start the notebook on existing embeddings."
      ],
      "metadata": {
        "id": "6soYhTrgauE8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sentence_transformers import util, SentenceTransformer\n",
        "\n",
        "embedding_model = SentenceTransformer(model_name_or_path=\"all-mpnet-base-v2\",\n",
        "                                      device=device)"
      ],
      "metadata": {
        "id": "bJTer6iU1ZSm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Perform a query using the dot product algorithm. It returns the top 5 results."
      ],
      "metadata": {
        "id": "mvqX-VbV1sx_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sentence_transformers import util\n",
        "\n",
        "query = \"blocked KV-cache\"\n",
        "print(f\"Query: {query}\")\n",
        "\n",
        "query_embedding = embedding_model.encode(query, convert_to_tensor=True)\n",
        "\n",
        "from time import perf_counter as timer\n",
        "\n",
        "start_time = timer()\n",
        "dot_scores = util.dot_score(a=query_embedding, b=embeddings)[0]\n",
        "end_time = timer()\n",
        "\n",
        "print(f\"Time take to get scores on {len(embeddings)} embeddings: {end_time-start_time:.5f} seconds.\")\n",
        "\n",
        "top_results_dot_product = torch.topk(dot_scores, k=5)\n",
        "top_results_dot_product"
      ],
      "metadata": {
        "id": "grvlmHce2RuI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a helper function to print wrapped text in the returned results in human readable format."
      ],
      "metadata": {
        "id": "FmY8y4Mg4lEt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import textwrap\n",
        "\n",
        "def print_wrapped(text, wrap_length=80):\n",
        "    wrapped_text = textwrap.fill(text, wrap_length)\n",
        "    print(wrapped_text)"
      ],
      "metadata": {
        "id": "6jnSrWiI4YRE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Loop through the top result tuple and match up the scores and indices and then use them to index the `pages_and_chunks` variable to get the relevant text chunk. Results are printed to the cell output in descending relevance score order."
      ],
      "metadata": {
        "id": "NhTaOPcF4ukf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"Query: '{query}'\\n\")\n",
        "print(\"Results:\")\n",
        "for score, idx in zip(top_results_dot_product[0], top_results_dot_product[1]):\n",
        "    print(f\"Score: {score:.4f}\")\n",
        "    print(\"Text:\")\n",
        "    print_wrapped(pages_and_chunks[idx][\"sentence_chunk\"])\n",
        "    print(f\"Page number: {pages_and_chunks[idx]['page_number']}\")\n",
        "    print(\"\\n\")"
      ],
      "metadata": {
        "id": "FhqO-7NK40Zd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Repeat the same query using cosine similarity."
      ],
      "metadata": {
        "id": "pNygOHaq5kik"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"blocked KV-cache\"\n",
        "print(f\"Query: {query}\")\n",
        "\n",
        "query_embedding = embedding_model.encode(query, convert_to_tensor=True)\n",
        "\n",
        "start_time = timer()\n",
        "cosine_scores = util.cos_sim(a=query_embedding, b=embeddings)[0]\n",
        "end_time = timer()\n",
        "\n",
        "print(f\"Time take to get scores on {len(embeddings)} embeddings: {end_time-start_time:.5f} seconds.\")\n",
        "\n",
        "top_results_cosine_product = torch.topk(cosine_scores, k=5)\n",
        "top_results_cosine_product"
      ],
      "metadata": {
        "id": "VznFZETA5sdd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generation"
      ],
      "metadata": {
        "id": "EjRd6U7i5s7F"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the llama-cpp-python package."
      ],
      "metadata": {
        "id": "xc8pnWLpugWR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install llama-cpp-python"
      ],
      "metadata": {
        "id": "B0zqh78D5vEF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the [microsoft/Phi-3-mini-4k-instruct-gguf](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct-gguf) model from the HF's Hub."
      ],
      "metadata": {
        "id": "QT3dL8_xuUTM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_cpp import Llama\n",
        "\n",
        "llm = Llama.from_pretrained(\n",
        "    repo_id=\"microsoft/Phi-3-mini-4k-instruct-gguf\",\n",
        "    filename=\"*-q4.gguf\",\n",
        "    verbose=False,\n",
        "    n_ctx=1024\n",
        ")"
      ],
      "metadata": {
        "id": "qtRSynRFZk59"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Provide a query and embed it."
      ],
      "metadata": {
        "id": "im7oBEuPupGk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"blocked KV-cache\"\n",
        "query_embedding = embedding_model.encode(query, convert_to_tensor=True)"
      ],
      "metadata": {
        "id": "KqX-CGspZrPV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Do the retrieval using cosine similarity."
      ],
      "metadata": {
        "id": "SSDr__hNutXK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cosine_scores = util.cos_sim(a=query_embedding, b=embeddings)[0]\n",
        "scores, indices = torch.topk(cosine_scores, k=2)\n",
        "context_items = [pages_and_chunks[i] for i in indices]"
      ],
      "metadata": {
        "id": "HEeuLOWWZv8Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the tokenizer associated to the selected generative model."
      ],
      "metadata": {
        "id": "iRRnVox5vBln"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"microsoft/Phi-3-mini-4k-instruct\")"
      ],
      "metadata": {
        "id": "oc0v0KAAhlft"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implement a custom function to format the prompt for the generator, by combining the query, the retrieved text chunks and a base prompt)."
      ],
      "metadata": {
        "id": "5Q-m4kvrvHWy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def prompt_formatter(query,\n",
        "                     context_items,\n",
        "                     tokenizer):\n",
        "    context = \"- \" + \"\\n- \".join([item[\"sentence_chunk\"] for item in context_items])\n",
        "\n",
        "    base_prompt = \"\"\"Based on the following context items, please answer the query.\n",
        "Make sure your answers are as explanatory as possible.\n",
        "\\nUse the following context items to answer the user query:\n",
        "{context}\n",
        "\\nRelevant passages: <extract relevant passages from the context here>\n",
        "User query: {query}\n",
        "Answer:\"\"\"\n",
        "\n",
        "    base_prompt = base_prompt.format(context=context, query=query)\n",
        "\n",
        "    dialogue_template = [\n",
        "        {\"role\": \"user\",\n",
        "        \"content\": base_prompt}\n",
        "    ]\n",
        "\n",
        "    prompt = tokenizer.apply_chat_template(conversation=dialogue_template,\n",
        "                                          tokenize=False,\n",
        "                                          add_generation_prompt=True)\n",
        "    return prompt"
      ],
      "metadata": {
        "id": "-0H24_1VZykT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Format the prompt."
      ],
      "metadata": {
        "id": "T3lyeV0rvj7r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = prompt_formatter(query=query,\n",
        "                          context_items=context_items,\n",
        "                          tokenizer=tokenizer)\n",
        "prompt"
      ],
      "metadata": {
        "id": "R2f-4AjMel-f"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Do the generation using the model in GGUF format."
      ],
      "metadata": {
        "id": "zFFP9E9MvnVD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output = llm(\n",
        "      prompt,\n",
        "      max_tokens=None,\n",
        "      stop=[\"Q:\", \"\\n\"],\n",
        "      echo=True\n",
        ")"
      ],
      "metadata": {
        "id": "q3u9kF4WZ8Po"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Display the answer."
      ],
      "metadata": {
        "id": "Mn9CgKW0v6m5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output"
      ],
      "metadata": {
        "id": "MtHZSn1RaGXK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The cells below require the local LanceDB database. Please go back to the [LanceDB section](#scrollTo=6RjcpGJimhv-) and execute the related code cells before moving further, if you haven't created and populated the database yet."
      ],
      "metadata": {
        "id": "D0io2EWrv46o"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Format the query for retrieval from the database."
      ],
      "metadata": {
        "id": "iu-rloNwww9J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_lancedb = prompt_formatter(query=query,\n",
        "                          context_items=results[0:2],\n",
        "                          tokenizer=tokenizer)\n",
        "prompt_lancedb"
      ],
      "metadata": {
        "id": "WnvGIblx9dE7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Do the generation using the model in GGUF format."
      ],
      "metadata": {
        "id": "HC8UErmow_VJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output = llm(\n",
        "      prompt_lancedb,\n",
        "      max_tokens=None,\n",
        "      stop=[\"Q:\", \"\\n\"],\n",
        "      echo=True\n",
        ")"
      ],
      "metadata": {
        "id": "uPM3a18292Cy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Display the answer."
      ],
      "metadata": {
        "id": "7ONPF0wgxJBi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output"
      ],
      "metadata": {
        "id": "TUlCXCKF_utt"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}